'''
Created on Aug 28, 2009

@author: mkiyer
'''

from matplotlib.patches import Polygon
import matplotlib.cm as cm
import matplotlib.pyplot as plt
import numpy as np
import operator

def plot_sample_coverage(chrom, start, end,
                         sample_groups, sample_coverages,
                         **kwargs):
    '''
    Plots individual sample coverage across the specified interval
    
    kwargs:
    axes: matplotlib Axes instance to plot into
    shade_start: start of region to shade
    shade_end: end of region to shade
    '''
    ax1 = kwargs.get('axes', plt.gca())     
    x = np.arange(0, end-start, 1)
    for sgroup in sample_groups:
        for s in sgroup.samples:
            ax1.plot(x, sample_coverages[s], color=sgroup.color, linewidth=1.0, label=s, alpha=0.75)
            ax1.annotate(s, xy=(x[0], np.max(sample_coverages[s])), xycoords='data',
                         size='xx-small', color=sgroup.color)

    shade_start = kwargs.get('shade_start', None)
    if shade_start is not None: 
        ax1.axvline(x=(shade_start - start), linewidth=2, color='black')
            
    ax1.grid(True)
    ax1.set_title("Individual samples", size='small')
    #ax1.set_title("%s:%d-%d" % (chrom, start, end), size='small')

def plot_sample_group_average_coverage(chrom, start, end,
                                       sample_groups, sample_coverages,
                                       **kwargs):
    '''
    Plots average coverage across the specified interval for each 
    sample group

    kwargs:
    axes: matplotlib Axes instance to plot into
    '''
    ax1 = kwargs.get('axes', plt.gca())        
    x = np.arange(0, end-start, 1)
    for sgroup in sample_groups:
        cov = np.zeros(end-start, dtype=np.float)
        for s in sgroup.samples:
            cov += sample_coverages[s]
        cov /= len(sgroup.samples)
        ax1.plot(x, cov, color=sgroup.color, linewidth=2.0, alpha=0.75)
        ax1.annotate(sgroup.name, xy=(x[0], np.max(cov)), xycoords='data',
                     size='xx-small', color=sgroup.color)

    shade_start = kwargs.get('shade_start', None)
    if shade_start is not None: 
        ax1.axvline(x=(shade_start - start), linewidth=2, color='black')
        
    ax1.grid(True)
    ax1.set_title('Sample groups', size='small')


def plot_boxplot(chrom, start, end,
                 sample_groups, sample_coverages,
                 **kwargs):
    '''
    sample_groups: list of SampleGroup objects.  the list order will be the
    order of the boxes    
    sample_coverages: dictionary keyed by sample name, values are 
    coverage data for that sample
    
    returns matplotlib Figure object
    
    kwargs:
    axes: matplotlib Axes instance to plot into
    title: title
    '''
    default_title = '%s:%d-%d' % (chrom, start, end)
    
    # construct boxplot data by making a list of vectors for
    # each sample group
    boxplotdata = [[sample_coverages[s] for s in sgroup.samples]
                   for sgroup in sample_groups]

    # setup plot
    ax1 = kwargs.get('axes', plt.gca())        

    # draw the boxes
    bp = ax1.boxplot(boxplotdata, notch=0, sym='+', vert=1, whis=1.5)
    plt.setp(bp['boxes'], color='black')
    plt.setp(bp['whiskers'], color='black')
    plt.setp(bp['fliers'], color='red', marker='+')

    # Add a horizontal grid to the plot, but make it very light in color
    # so we can use it for reading data values but not be distracting
    ax1.yaxis.grid(True, linestyle='-', which='major', color='lightgrey',
                  alpha=0.5)
    # Hide these grid behind plot objects
    ax1.set_axisbelow(True)
    
    # Set plot labels
    ax1.set_title('Sample groups', size='small')
    #ax1.set_title(kwargs.get('title', default_title))
    #ax1.set_xlabel('Sample group', size='x-small')
    ax1.set_ylabel('Coverage (RPKM)', size='small')

    # Now fill the boxes with desired colors
    boxColors = [sgroup.color for sgroup in sample_groups]
    numBoxes = len(boxplotdata)
    medians = range(numBoxes)
    for i in range(numBoxes):
        box = bp['boxes'][i]
        boxX = []
        boxY = []
        for j in range(5):
            boxX.append(box.get_xdata()[j])
            boxY.append(box.get_ydata()[j])
        boxCoords = zip(boxX, boxY)
        boxPolygon = Polygon(boxCoords, facecolor=boxColors[i])
        ax1.add_patch(boxPolygon)
        # Now draw the median lines back over what we just filled in
        med = bp['medians'][i]
        medianX = []
        medianY = []
        for j in range(2):
            medianX.append(med.get_xdata()[j])
            medianY.append(med.get_ydata()[j])
            ax1.plot(medianX, medianY, 'k')
            medians[i] = medianY[0]
        # Finally, overplot the sample averages, with horizontal alignment
        # in the center of each box
        ax1.plot([np.average(med.get_xdata())], [np.average(boxplotdata[i])],
                 color='w', marker='*', markeredgecolor='k')

    # Set the axes ranges and axes labels
    ax1.set_xlim(0.5, numBoxes+0.5)
    # TODO: setting ylim appears to screw things up
    bottom, top = ax1.get_ylim()
    #top = 1000
    #bottom = -5
    #ax1.set_ylim(bottom, top)
    labels = [sgroup.name for sgroup in sample_groups]
    xtickNames = plt.setp(ax1, xticklabels=labels)
    #plt.setp(xtickNames, rotation=45, fontsize=8)
    plt.setp(xtickNames, rotation=0, size='x-small')
    
    # Due to the Y-axis scale being different across samples, it can be
    # hard to compare differences in medians across the samples. Add upper
    # X-axis tick labels with the sample medians to aid in comparison
    # (just use two decimal places of precision)
    pos = np.arange(numBoxes)+1
    upperLabels = [str(np.round(s, 2)) for s in medians]
    weights = ['bold', 'semibold']
    for tick,label in zip(range(numBoxes),ax1.get_xticklabels()):
       k = tick % 2
       ax1.text(pos[tick], top-(top*0.05), upperLabels[tick],
            horizontalalignment='center', size='x-small', weight=weights[k],
            color=boxColors[tick])
    
    # Finally, add a basic legend
    ystart = 0.045
    yincrement = 0.02
    for sgroup in sample_groups:
        plt.figtext(0.80, ystart, sgroup.name,
                    backgroundcolor=sgroup.color, color='black', weight='roman',
                    size='x-small')
        ystart += yincrement
    plt.figtext(0.80, 0.015, '*', color='white', backgroundcolor='silver',
               weight='roman', size='medium')
    plt.figtext(0.815, 0.013, ' Average Value', color='black', weight='roman',
               size='x-small')


def plot_barplot(chrom, start, end,
                 sample_groups, sample_coverages,
                 **kwargs):
    '''
    sample_groups: list of SampleGroup objects.  the list order will be the
    order of the bar groups    
    sample_coverages: dictionary keyed by sample name, values are 
    coverage data for that sample
    
    returns matplotlib Figure object
    
    kwargs:
    axes: matplotlib Axes instance to plot into
    title: title
    '''
    default_title = '%s:%d-%d' % (chrom, start, end)
    # setup plot
    ax = kwargs.get('axes', plt.gca())    
    # construct boxplot data by making a list of vectors for
    # each sample group
    barlabels = []
    bardata = []
    barcolors = []
    for sgroup in sample_groups:
        barsamples = [(s, sample_coverages[s]) for s in sgroup.samples]
        barsamples = sorted(barsamples, key=operator.itemgetter(1), reverse=False)
        barlabels.extend([x[0] for x in barsamples])
        bardata.extend([x[1] for x in barsamples])
        barcolors.extend([sgroup.color for x in barsamples])

    ind = np.arange(len(bardata))
    width = 0.8
    ax.bar(ind, bardata, width, color=barcolors)
    # labels
    ax.set_ylabel('Coverage (RPKM)')
    ax.set_title("Individual samples", size='x-small')

    #ax.set_title('Scores by group and gender')
    ax.set_xticks(ind+width/2.0)
    ax.set_xticklabels(barlabels, rotation='vertical', size='xx-small')
    #ax.legend( (rects1[0], rects2[0]), ('Men', 'Women') )
    #autolabel(rects1)
    #autolabel(rects2)

#    def autolabel(rects):
#        # attach some text labels
#        for rect in rects:
#            height = rect.get_height()
#            ax.text(rect.get_x()+rect.get_width()/2., 1.05*height, '%d'%int(height),
#                    ha='center', va='bottom')



